#include "modules/drivers/orin_gnss_v2/optimizer/kfpose/kfpose.h"

DEFINE_string(gnss_best_pose_opt_topic, "/apollo/sensor/gnss/best_pose_opt",
              "Optimized GNSS BestPose topic name");

namespace apollo {
namespace drivers {
namespace gnss {

bool KFPoseComponent::Init() {
  AINFO << "KFPoseComponent component init";
  Config gnss_config;
  if (!apollo::cyber::common::GetProtoFromFile(config_file_path_,
                                               &gnss_config)) {
    AERROR << "Unable to load gnss conf file: " << config_file_path_;
    return false;
  }
  wgs84pj_source_ = pj_init_plus("+proj=latlong +ellps=WGS84");
  utm_target_ = pj_init_plus(gnss_config.proj4_text().c_str());

  gnssbestpose_writer_ =
      node_->CreateWriter<GnssBestPose>(FLAGS_gnss_best_pose_opt_topic);
  bestpose_reader_ =
      node_->CreateReader<GnssBestPose>(FLAGS_gnss_best_pose_topic, nullptr);
  return true;
}

bool KFPoseComponent::Proc(const std::shared_ptr<Gnss>& gnss) {
  auto gnss_msg = gnss.get();

  double x = gnss_msg->position().lon();
  double y = gnss_msg->position().lat();
  double z = gnss_msg->position().height();

  double std_x = gnss_msg->position_std_dev().x();
  double std_y = gnss_msg->position_std_dev().y();
  double std_z = gnss_msg->position_std_dev().z();

  double vel_x = gnss_msg->linear_velocity().x();
  double vel_y = gnss_msg->linear_velocity().y();
  double vel_z = gnss_msg->linear_velocity().z();

  x *= DEG_TO_RAD;
  y *= DEG_TO_RAD;

  pj_transform(wgs84pj_source_, utm_target_, 1, 1, &x, &y, NULL);

  if (!gnss_kf_.get()) {
    Eigen::Matrix<double, 6, 1> kf_x;
    Eigen::Matrix<double, 6, 6> kf_P;
    Eigen::Matrix<double, 6, 6> kf_F;
    Eigen::Matrix<double, 6, 6> kf_H;

    kf_x << x, y, z, 0.0, 0.0, 0.0;

    Eigen::VectorXd diag_P(6);
    diag_P << std_x * std_x, std_y * std_y, std_z * std_z, 0.01, 0.01, 0.01;
    kf_P.setZero();
    kf_P.diagonal() = diag_P;

    gnss_kf_.reset(new KalmanFilter<double, 6, 6, 0>(kf_x, kf_P));

    kf_F.setIdentity();
    kf_F.block<3, 3>(0, 3) = Eigen::Matrix<double, 3, 3>::Identity();
    gnss_kf_->SetTransitionMatrix(kf_F);

    Eigen::Matrix<double, 6, 6> kf_Q;
    kf_Q.setZero();
    kf_Q.diagonal() = diag_P;
    gnss_kf_->SetTransitionNoise(kf_Q);

    Eigen::Matrix<double, 6, 6> kf_R;
    kf_R.setZero();
    kf_R.diagonal() = diag_P;
    gnss_kf_->SetObservationNoise(kf_R);

    kf_H.setIdentity();
    gnss_kf_->SetObservationMatrix(kf_H);
  } else {
    gnss_kf_->Predict();
    Eigen::Matrix<double, 6, 1> kf_z;
    Eigen::Matrix<double, 6, 6> kf_R;
    kf_z << x, y, z, vel_x, vel_y, vel_z;

    Eigen::VectorXd diag_R(6);
    diag_R << std_x * std_x, std_y * std_y, std_z * std_z, 0.01, 0.01, 0.01;
    kf_R.setZero();
    kf_R.diagonal() = diag_R;

    // gnss_kf_->SetTransitionNoise(kf_R);
    gnss_kf_->SetObservationNoise(kf_R);
    gnss_kf_->Correct(kf_z);

    auto kf_x_result = gnss_kf_->GetStateEstimate();
    auto kf_P_result = gnss_kf_->GetStateCovariance();

    double x_kf = kf_x_result(0);
    double y_kf = kf_x_result(1);
    double z_kf = kf_x_result(2);

    double std_x_kf = sqrt(kf_P_result(0, 0));
    double std_y_kf = sqrt(kf_P_result(1, 1));
    double std_z_kf = sqrt(kf_P_result(2, 2));

    pj_transform(utm_target_, wgs84pj_source_, 1, 1, &x_kf, &y_kf, NULL);
    x_kf *= RAD_TO_DEG;
    y_kf *= RAD_TO_DEG;

    bestpose_reader_->Observe();
    if (bestpose_reader_->Empty()) {
      AERROR << "no bestpose";
      return true;
    }
    auto bestpos_ptr = bestpose_reader_->GetLatestObserved();

    if (bestpos_ptr != nullptr) {
      auto bestpos = std::make_shared<GnssBestPose>();
      bestpos->CopyFrom(*bestpos_ptr.get());
      auto header = bestpos->mutable_header();
      header->set_timestamp_sec(gnss_msg->header().timestamp_sec());
      header->set_sequence_num(gnss_msg->header().sequence_num());

      bestpos->set_latitude(y_kf);
      bestpos->set_longitude(x_kf);
      bestpos->set_height_msl(z_kf);
      bestpos->set_undulation(0.0);  // gnss msg already add undulation
      bestpos->set_latitude_std_dev(std_y_kf);
      bestpos->set_longitude_std_dev(std_x_kf);
      bestpos->set_height_std_dev(std_z_kf);
      gnssbestpose_writer_->Write(bestpos);
    }
  }

  return true;
}
}  // namespace gnss
}  // namespace drivers
}  // namespace apollo